#include "CurveFitter.h"

using namespace cv;
using namespace std;

#define PI (3.1415926)
#define DEBUG

sinLeastSquare::sinLeastSquare(deque<Point2f> dataV, double omega): cv::MinProblemSolver::Function()
{
    this->dataV = dataV;
    this->omega = omega;
}

int sinLeastSquare::getDims() const{ return 2; }
void sinLeastSquare::getGradient(const double* arg,double* grad)
{
    grad[0] = grad[1] = 0;
    //grad[2] = 0.;
    for(int i = 0; i < dataV.size(); i++)
    {
        // double x = (*ptr_dataV)[i].x;
        // double y = (*ptr_dataV)[i].y;
        double x = dataV[i].x;
        double y = dataV[i].y;

        double f = arg[0] * sin(omega * x + arg[1]) - y;
        grad[0] += 2 * f * sin(omega * x + arg[1]); 
        grad[1] += 2 * f * cos(omega * x + arg[1]) * arg[0];
        // grad[2] += 2 * f * cos(omega * x + arg[1]) * arg[0] * x ;
    }
}
double sinLeastSquare::calc(const double* arg)const
    {
        double sum = 0;
        for(int i = 0; i < dataV.size(); i++)
        {
            double x = dataV[i].x;
            double y = dataV[i].y;
            double s = (arg[0] * sin(omega * x + arg[1]) - y);
            sum += s*s;
            return sum; 
        }
    }   


/*
大能量机关转速按照三角函数呈周期性变化。速度目标函数为:spd=0.785*sin(1.884*t)+1.305，其中 spd
的单位为 rad/s，t 的单位为 s，且每次大能量机关进入可激活状态时，t 重置为零。 

大能量机关的实际转速与速度目标函数的时间误差在 500ms 内。
*/
CurveFitter::CurveFitter(const int fitLineNumThreshold,
                         const int fitLineNumUpperBound,
                         const int fitSinNumThreshold,
                         const int fitSinNumUpperBound)
{
    this->powerRuneCircleCenter = Point2f(0.,0.);
    this->fitLineNumThreshold = fitLineNumThreshold;
    this->fitLineNumUpperBound = fitLineNumUpperBound;
    this->curveFitDataV = this->fitSinDataV =  deque<Point2f>();
    this->base_angle = -1.0; 
    this->startTime = getTime_ms();
    this->fitLineK = 0.;
    this->solver=cv::ConjGradSolver::create();
    this->sinA = -0.785 / 1.884;
    this->sinB = 0.; 
    this->sinOmega = 1.884;

    this->fitSinNumThreshold = fitSinNumThreshold;
    this->fitSinNumUpperBound = fitSinNumUpperBound;

    cv::TermCriteria termCriteria;
    termCriteria.type = TermCriteria::MAX_ITER;
    termCriteria.epsilon = 0.001;
    termCriteria.maxCount = 100;

    this->solver->setTermCriteria(termCriteria);
}

void CurveFitter::initSet(const Point2f powerRuneCircleCenter, const double baseAngle)
{
    this->powerRuneCircleCenter = powerRuneCircleCenter;
    this->base_angle = baseAngle;
    this->startTime = getTime_ms();
}

void CurveFitter::inputCurveFitDataV(deque<Point2f> &circleDataV)
{
    this->curveFitDataV = circleDataV;
    this->fitSinDataV = circleDataV; 
}

bool CurveFitter::isFirstSetBaseAngle() const
{
    if(this->base_angle < 0)
        return true;
    else 
        return false; 
}

//https://blog.csdn.net/guduruyu/article/details/69505487
void CurveFitter::fitLine()
{
    vector<Point2f> lineFitDataV;
    for(int i = 0; i < this->curveFitDataV.size(); i++)
    {
        lineFitDataV.push_back(this->curveFitDataV[i]); //FIXME: this is stupid 
    }

	cv::fitLine(lineFitDataV, this->fitLinePara, cv::DIST_L2, 0, 1e-2, 1e-2); // use least sqaure method to fit line

}

// \phi = A * sin(\Omega * t + B)
void CurveFitter::fitSin(deque<Point2f> dataQue)
{
    double sum = 0.;
    sinMean = 0.;
    for(int i = 0; i < dataQue.size(); i++)
    {
        double t_s = dataQue[i].x;
        dataQue[i].y -= this->fitLineK * t_s;  // 修改dataQue
        sum+= dataQue[i].y;
    }
    sinMean = sum / dataQue.size();
#ifdef DEBUG
    // printf("mean: %lf\n", sinMean); 
#endif
    for(int i = 0; i < dataQue.size(); i++)
    {
        dataQue[i].y -= sinMean;
    }

    this->ptr_F = cv::Ptr<cv::MinProblemSolver::Function>(new sinLeastSquare(dataQue, sinOmega));
    solver->setFunction(ptr_F);
    Mat arg = (cv::Mat_<double>(2,1)<< this->sinA,this->sinB);
    double min_result=solver->minimize(arg);
    
    //update sinA and sinB
    this->sinA = arg.at<double>(0,0);
    this->sinB = arg.at<double>(1,0);
    // this->sinOmega = arg.at<double>(2,0);

#ifdef DEBUG
    // std::cout<<"min_result:\n"<<min_result<<std::endl;
#endif
}

/*
    坐标系
    |------> x (左上角为(0,0))
    |
    v
    y
*/

// 输入的时间单位是s！！！
double CurveFitter::getFitAngle(double t_s)
{
    double a = t_s * fitLineK + sinA * sin(sinOmega * t_s + sinB);
    return fmod(a, 2*PI);
}

Point2f CurveFitter::predictNewArmorBoxCenter(const double delta_t_ms, const Point2f currentArmorBoxCenter, const double currentAngle,
                                            const double circleRadius, const Point2f inputBias)
{
    double current_t_ms = getTime_ms() - this->startTime;
    double current_t_s = current_t_ms / 1000.; 
    double t_s = (current_t_ms + delta_t_ms) / 1000.; 

    // double newAngle = fmod((-t_s * fitLineK), 2 * PI); 
    double newAngle = getFitAngle(t_s);
    double deltaAngle = fmod(newAngle - getFitAngle(current_t_s), 2*PI); 

    double referenceAngle = 0.;
    if(currentArmorBoxCenter.x - powerRuneCircleCenter.x == 0 && currentArmorBoxCenter.y - powerRuneCircleCenter.y == 0 )
        referenceAngle = 0.;
    else 
        referenceAngle = atan2(
            currentArmorBoxCenter.x - powerRuneCircleCenter.x,
            currentArmorBoxCenter.y - powerRuneCircleCenter.y
                            );
    
    if(referenceAngle < 0.)
    {
        referenceAngle += 2*PI; //把角度纳入统一的一个极坐标系下考虑
        /*
                    | PI
                    |
         3PI/2------|------ PI / 2
                    |
                2PI | 0
        */
    }
    // TODO: add powerRune rotate direction judge
    double targetAngle = fmod(referenceAngle + deltaAngle, 2 * PI);

    double newY = cos(targetAngle) * circleRadius; 
    double newX = sin(targetAngle) * circleRadius; 
    Point2f targetArmorBox = Point2f(newX, newY) + this->powerRuneCircleCenter;

    targetArmorBox += inputBias;
    
#ifdef DEBUG
    // printf("newAngle = %lf, referenceAngle = %lf, deltaAngle = %lf\n",newAngle, referenceAngle, deltaAngle);
    // printf("===============[Predict]===============\n");
    // printf("[current_t_ms = %lf, next_t_s = %lf]\n", current_t_ms, t_s);
    // printf("newAngle = %lf, \ndeltaAngle = %lf, \nreferenceAngle = %lf,\ntargetAngle = %lf \n\nnewX = %lf, newY = %lf\n", newAngle, deltaAngle, referenceAngle, targetAngle, newX, newY);
#endif
    return targetArmorBox;
}

void CurveFitter::curveFitting(Point2f armorBoxCenter, Point2f& predictArmorBoxCenter, const double circleRadius, const Point2f inputBias, const double delta_t_ms)
{
    static Mat canvas;
    static double lastArmorBoxCenterAngle = 0;
    static int loopCount = 0;  
    // 1. first get angle of armorBoxCenter & get time : [t, \theta]

    // assert armorBoxCenterAngle in range -pi - pi 
    double armorBoxCenterAngle;

    if((armorBoxCenter.y - powerRuneCircleCenter.y) == 0 &&
       (armorBoxCenter.x - powerRuneCircleCenter.x) == 0) //atan2 domain error 
        armorBoxCenterAngle = 0.; 
    else 
        armorBoxCenterAngle = atan2(armorBoxCenter.x - powerRuneCircleCenter.x,
            armorBoxCenter.y - powerRuneCircleCenter.y
                                      );

    if(armorBoxCenterAngle < 0.)
    {
        armorBoxCenterAngle += 2 * PI; // make sure armorBoxCenterAngle in range [0, 2PI]
        /*
                    | PI
                    |
         3PI/2------|------ PI / 2
                    |
                2PI | 0
        */
    }
    ///// Error Handling 
    if(lastArmorBoxCenterAngle == 0)
        lastArmorBoxCenterAngle = armorBoxCenterAngle;
    else if(lastArmorBoxCenterAngle - armorBoxCenterAngle > PI)
    {
        loopCount -= 1;
        // printf("loop: %d\n\n", loopCount);
    }
    else if(lastArmorBoxCenterAngle - armorBoxCenterAngle < -PI)
    {
        loopCount += 1;
    }
    lastArmorBoxCenterAngle = armorBoxCenterAngle;
    //////
    
    double t_s = (getTime_ms() - this->startTime) / 1000.;

    // printf("real angle %lf\n",armorBoxCenterAngle);

    double recordAngle = armorBoxCenterAngle - loopCount * 2 * PI; 

    this->curveFitDataV.push_back(Point2f(t_s, recordAngle));
    this->fitSinDataV.push_back(Point2f(t_s, recordAngle)); //提前把data放到fitSinDataV中

    chrono::steady_clock::time_point t1, t2; 

    // 2. if curveFitDataV.size() > threshold, should start to fit curve 
    if(this->curveFitDataV.size() >= fitLineNumThreshold)
    {
        // fit curve
        // 2.1. Use least square to fit k 
        t1 = chrono::steady_clock::now();
        fitLine();
        t2 = chrono::steady_clock::now();
        cout << "curveFitter: fitLine time: " << (static_cast<chrono::duration<double, std::milli>>(t2 - t1)).count() << " ms" << endl;
        
        if(this->fitLinePara[1] == 0.0)
            this->fitLineK = 0.;
        else 
            this->fitLineK = this->fitLinePara[0] / this->fitLinePara[1]; // 负号根据实际调整加上的
        
        // 2.2 Use least square to fit sin wave
        if(this->fitSinDataV.size() >= fitSinNumThreshold) // 其实这个if不太必要了，在现在的写法下外层if满足时内层if也应该是满足的
        {
            t1 = chrono::steady_clock::now();
            fitSin(this->fitSinDataV);
            t2 = chrono::steady_clock::now();
            cout << "curveFitter: fitSin time: " << (static_cast<chrono::duration<double, std::milli>>(t2 - t1)).count() << " ms" << endl;
        }
    }
    if(this->curveFitDataV.size() > fitLineNumUpperBound)
        this->curveFitDataV.pop_front();
    if(this->fitSinDataV.size() > fitSinNumUpperBound)
        this->fitSinDataV.pop_front();

    // 3. after fitting curve, give prediction based on new timestamp 
    t1 = chrono::steady_clock::now();
    predictArmorBoxCenter = predictNewArmorBoxCenter(delta_t_ms, armorBoxCenter, recordAngle, circleRadius, inputBias);
    t2 = chrono::steady_clock::now();
    cout << "curveFitter: predictArmorBoxCenter time: " << (static_cast<chrono::duration<double, std::milli>>(t2 - t1)).count() << " ms" << endl;

#ifdef DEBUG
    // printf("armorBoxCenter x: %lf, y: %lf\n", armorBoxCenter.x, armorBoxCenter.y);
    // printf("predictArmorBoxCenter x: %lf, y: %lf\n", predictArmorBoxCenter.x, predictArmorBoxCenter.y);
    // printf("loop: %d\n", loopCount);
    // curveFittingDebug(canvas); // For debugging
#endif 

    return; 
}

void CurveFitter::setPowerRuneCircleCenter(const Point2f newCircleCenter)
{
    this->powerRuneCircleCenter = newCircleCenter; 
    this->fitSinDataV.clear(); //TODO
    this->curveFitDataV.clear();
}

Point2f CurveFitter::getPowerRuneCircleCenter() const 
{
    return this->powerRuneCircleCenter;
}

void CurveFitter::curveFittingDebug(Mat& canvas)
{
    static const bool drawArmorBoxCenterAngleCurve = true; 

    if(drawArmorBoxCenterAngleCurve)
    {
        const int img_size = 5000; 
        canvas = cv::Mat(1500, img_size, CV_8UC3, Scalar(0, 0, 0));
        
        for(int i = 0; i < this->curveFitDataV.size(); i++)
        {
            double t_s = this->curveFitDataV[i].x; 
            double angle = this->curveFitDataV[i].y;
            double afterLineCorrect_angle = angle - this->fitLineK * t_s;  

            double sinPredict_angle = sinA * sin(sinOmega * t_s + sinB); //正负号需要实际测试的时候确定
            double afterSinCorrect_angle = afterLineCorrect_angle - sinPredict_angle - this->sinMean;

            if(i == this->curveFitDataV.size()-1 )
            {
                // printf("===============[Summary]===============\n");
                // printf("time: %lf, originAngle: %lf\n", t_s, angle);
                // printf("time: %lf, after line correct Angle: %lf\n", t_s, afterLineCorrect_angle);
                // printf("time: %lf, after sin correct Angle: %lf\n", t_s, afterSinCorrect_angle);
                // printf("K: %lf\n", this->fitLineK);
                // printf("sinA: %lf, sinOmega: %lf, sinB: %lf\n", sinA, sinOmega, sinB);
                // printf("\n");
            }
            
            // Scale the data for drawing 
            double x = fmod((t_s * 1000. / 4.0), img_size);  
            double y_afterLine = ( (int)((afterLineCorrect_angle + 7) * 100)); 
            double y_afterSinCorrect = ( (int)((afterSinCorrect_angle+ 7) * 100)); 
            double y_sinFit = ( (int)((sinPredict_angle + 7) * 100)); 

            circle(canvas, Point2f(x , y_afterLine), 1, Scalar(0, 0, 255), 30);  // red : afterlineCorrect
            circle(canvas, Point2f(x , y_afterSinCorrect), 1, Scalar(0, 255, 0), 30);  // green: afterSin
            circle(canvas, Point2f(x , y_sinFit), 1, Scalar(255, 0, 255), 30); //purple: sinPrecit
        }
        namedWindow("ArmorBoxCenterAngleCurve", WINDOW_NORMAL);
        imshow("ArmorBoxCenterAngleCurve", canvas);
        waitKey(1);
    }
}
